"""Ingest local text/markdown files into vector store.

可直接脚本方式运行 (python rag/ingest.py <路径>) 或模块方式 (python -m rag.ingest <路径>)。
修复: 绝对导入避免 relative import 错误。
"""
from pathlib import Path
from typing import Iterable, List

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.docstore.document import Document

from config import config

SUPPORTED_SUFFIX = {".txt", ".md"}


def iter_files(input_path: Path) -> Iterable[Path]:
    if input_path.is_file():
        yield input_path
    else:
        for p in input_path.rglob("*"):
            if p.suffix.lower() in SUPPORTED_SUFFIX:
                yield p


def load_documents(paths: List[Path]) -> List[Document]:
    docs: List[Document] = []
    for p in paths:
        try:
            text = p.read_text(encoding="utf-8", errors="ignore")
        except Exception as e:
            print(f"Skip {p}: {e}")
            continue
        docs.append(Document(page_content=text, metadata={"source": str(p)}))
    return docs


def chunk_documents(docs: List[Document]):
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=config.chunk_size, chunk_overlap=config.chunk_overlap)
    return splitter.split_documents(docs)


def build_vector_store(chunks: List[Document]):
    embeddings = HuggingFaceEmbeddings(model_name=config.embedding_model, encode_kwargs={"batch_size": config.batch_size})
    vs = Chroma.from_documents(chunks, embedding=embeddings,
                               collection_name=config.collection_name, persist_directory=str(config.persist_dir))
    vs.persist()
    return vs


def ingest(path: str):
    input_path = Path(path)
    all_files = list(iter_files(input_path))
    if not all_files:
        raise SystemExit("No supported files found to ingest.")
    docs = load_documents(all_files)
    chunks = chunk_documents(docs)
    build_vector_store(chunks)
    print(
        f"Ingested {len(docs)} docs -> {len(chunks)} chunks into {config.persist_dir}")


if __name__ == "__main__":
    ingest("C:\\planner")
